import numpy as np
import yaml
import spatialmath as sm
import pinocchio as pin

from ..robot_state import RobotState


class IkRes:

    def __init__(self):
        super().__init__()

        self.status = 0
        self.itr = 0
        self.err = np.array([])
        self.jointPosRes = np.array([])


class Robotics:

    def __init__(self, urdf_path: str):
        super().__init__()

        root_joint = pin.JointModelFreeFlyer()
        self.model_biped = pin.buildModelFromUrdf(urdf_path, root_joint)
        self.model_biped_fixed = pin.buildModelFromUrdf(urdf_path)
        self.data_biped = self.model_biped.createData()
        self.data_biped_fixed = self.model_biped_fixed.createData()
        self.model_nv = self.model_biped.nv
        self.J_l = np.zeros((6, self.model_nv))
        self.J_r = np.zeros((6, self.model_nv))
        self.J_base = np.zeros((6, self.model_nv))
        self.J_hip_link = np.zeros((6, self.model_nv))
        self.dJ_l = np.zeros((6, self.model_nv))
        self.dJ_r = np.zeros((6, self.model_nv))
        self.dJ_base = np.zeros((6, self.model_nv))
        self.q = np.zeros(self.model_biped.nq)
        self.dq = np.zeros(self.model_nv)
        self.ddq = np.zeros(self.model_nv)
        self.Rcur = np.eye(3)
        self.dyn_M = np.zeros((self.model_nv, self.model_nv))
        self.dyn_M_inv = np.zeros((self.model_nv, self.model_nv))
        self.dyn_C = np.zeros((self.model_nv, self.model_nv))
        self.dyn_Ag = np.zeros((6, self.model_nv))
        self.dyn_dAg = np.zeros((6, self.model_nv))
        self.dyn_G = np.zeros(self.model_nv)
        self.dyn_Non = np.zeros(self.model_nv)

        self.l_ankle_joint = self.model_biped.getJointId("leg_l6_joint")
        self.r_ankle_joint = self.model_biped.getJointId("leg_r6_joint")
        self.l_hip_joint = self.model_biped.getJointId("leg_l1_joint")
        self.r_hip_joint = self.model_biped.getJointId("leg_r1_joint")
        self.l_ankle_joint_fixed = self.model_biped_fixed.getJointId("leg_l6_joint")
        self.r_ankle_joint_fixed = self.model_biped_fixed.getJointId("leg_r6_joint")
        self.l_hip_joint_fixed = self.model_biped_fixed.getJointId("leg_l1_joint")
        self.r_hip_joint_fixed = self.model_biped_fixed.getJointId("leg_r1_joint")
        self.base_joint = self.model_biped.getJointId("root_joint")

        self.Jcom = np.zeros((3, self.model_nv))
        self.fe_r_pos = np.zeros(3)
        self.fe_l_pos = np.zeros(3)
        self.base_pos = np.zeros(3)
        self.fe_r_pos_body = np.zeros(3)
        self.fe_l_pos_body = np.zeros(3)
        self.hip_r_pos = np.zeros(3)
        self.hip_l_pos = np.zeros(3)
        self.hip_link_pos = np.zeros(3)
        self.hip_r_pos_body = np.zeros(3)
        self.hip_l_pos_body = np.zeros(3)
        self.hip_link_rot = np.zeros((3, 3))
        self.fe_r_rot = np.zeros((3, 3))
        self.fe_l_rot = np.zeros((3, 3))
        self.base_rot = np.zeros((3, 3))
        self.fe_r_rot_body = np.zeros((3, 3))
        self.fe_l_rot_body = np.zeros((3, 3))
        self.CoM_pos = np.zeros(3)
        self.inertia = np.zeros((3, 3))

    def read_state(self, robot_state: RobotState):
        self.q[:] = robot_state.q
        self.dq[:] = robot_state.dq
        self.dq[:3] = robot_state.base_rot.T @ self.dq[:3]
        self.dq[3: 6] = robot_state.base_rot.T @ self.dq[3: 6]
        self.ddq[:] = robot_state.ddq

    def write_state(self, robot_state: RobotState):
        robot_state.J_l[:, :] = self.J_l
        robot_state.J_r[:, :] = self.J_r
        robot_state.J_base[:, :] = self.J_base
        robot_state.dJ_l[:, :] = self.dJ_l
        robot_state.dJ_r[:, :] = self.dJ_r
        robot_state.dJ_base[:, :] = self.dJ_base
        robot_state.J_hip_link[:, :] = self.J_hip_link
        robot_state.fe_l_pos_W[:] = self.fe_l_pos
        robot_state.fe_r_pos_W[:] = self.fe_r_pos
        robot_state.fe_l_pos_L[:] = self.fe_l_pos_body
        robot_state.fe_r_pos_L[:] = self.fe_r_pos_body
        robot_state.fe_l_rot_W[:, :] = self.fe_l_rot
        robot_state.fe_r_rot_W[:, :] = self.fe_r_rot
        robot_state.fe_l_rot_L[:, :] = self.fe_l_rot_body
        robot_state.fe_r_rot_L[:, :] = self.fe_r_rot_body
        robot_state.hip_l_pos_L[:] = self.hip_l_pos_body
        robot_state.hip_r_pos_L[:] = self.hip_r_pos_body
        robot_state.hip_l_pos_W[:] = self.hip_l_pos
        robot_state.hip_r_pos_W[:] = self.hip_r_pos
        robot_state.hip_link_pos[:] = self.hip_link_pos
        robot_state.hip_link_rot[:, :] = self.hip_link_rot

        robot_state.dyn_M[:, :] = self.dyn_M
        robot_state.dyn_M_inv[:, :] = self.dyn_M_inv
        robot_state.dyn_C[:, :] = self.dyn_C
        robot_state.dyn_G[:] = self.dyn_G
        robot_state.dyn_Ag[:] = self.dyn_Ag
        robot_state.dyn_dAg[:] = self.dyn_dAg
        robot_state.dyn_Non[:] = self.dyn_Non

        robot_state.pCoM_W[:] = self.CoM_pos
        robot_state.Jcom_W[:, :] = self.Jcom

        robot_state.inertia[:, :] = self.inertia

    def computeJ_dJ(self):
        pin.forwardKinematics(self.model_biped, self.data_biped, self.q)
        pin.jacobianCenterOfMass(self.model_biped, self.data_biped, self.q, True)
        pin.computeJointJacobiansTimeVariation(self.model_biped, self.data_biped, self.q, self.dq)
        pin.updateGlobalPlacements(self.model_biped, self.data_biped)
        self.J_l[:, :] = pin.getJointJacobian(self.model_biped, self.data_biped, self.l_ankle_joint,
                                              pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
        self.J_r[:, :] = pin.getJointJacobian(self.model_biped, self.data_biped, self.r_ankle_joint,
                                              pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
        self.J_base[:, :] = pin.getJointJacobian(self.model_biped, self.data_biped, self.base_joint,
                                                 pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
        self.J_hip_link[:, :] = pin.getJointJacobian(self.model_biped, self.data_biped, self.base_joint,
                                                     pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)

        self.dJ_l[:, :] = pin.getJointJacobianTimeVariation(self.model_biped, self.data_biped, self.l_ankle_joint,
                                                            pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
        self.dJ_r[:, :] = pin.getJointJacobianTimeVariation(self.model_biped, self.data_biped, self.r_ankle_joint,
                                                            pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)
        self.dJ_base[:, :] = pin.getJointJacobianTimeVariation(self.model_biped, self.data_biped, self.base_joint,
                                                               pin.ReferenceFrame.LOCAL_WORLD_ALIGNED)

        self.fe_l_pos[:] = self.data_biped.oMi[self.l_ankle_joint].translation
        self.fe_l_rot[:, :] = self.data_biped.oMi[self.l_ankle_joint].rotation
        self.fe_r_pos[:] = self.data_biped.oMi[self.r_ankle_joint].translation
        self.fe_r_rot[:, :] = self.data_biped.oMi[self.r_ankle_joint].rotation
        self.hip_l_pos[:] = self.data_biped.oMi[self.l_hip_joint].translation
        self.hip_r_pos[:] = self.data_biped.oMi[self.r_hip_joint].translation
        self.base_pos[:] = self.data_biped.oMi[self.base_joint].translation
        self.base_rot[:, :] = self.data_biped.oMi[self.base_joint].rotation
        self.hip_link_rot[:, :] = self.data_biped.oMi[self.base_joint].rotation
        self.Jcom[:, :] = self.data_biped.Jcom

        Mpj = np.eye(self.model_nv)
        Mpj[0: 3, 0: 3] = self.base_rot.T
        Mpj[3: 6, 3: 6] = self.base_rot.T
        self.J_l[:, :] = self.J_l @ Mpj
        self.J_r[:, :] = self.J_r @ Mpj
        self.J_base[:, :] = self.J_base @ Mpj
        self.dJ_l[:, :] = self.dJ_l @ Mpj
        self.dJ_r[:, :] = self.dJ_r @ Mpj
        self.dJ_base[:, :] = self.dJ_base @ Mpj
        self.J_hip_link[:, :] = self.J_hip_link @ Mpj
        self.Jcom[:, :] = self.Jcom @ Mpj

        q_fixed = self.q[7: 7 + self.model_biped_fixed.nv]
        pin.forwardKinematics(self.model_biped_fixed, self.data_biped_fixed, q_fixed)
        pin.updateGlobalPlacements(self.model_biped_fixed, self.data_biped_fixed)
        self.fe_l_pos_body[:] = self.data_biped_fixed.oMi[self.l_ankle_joint_fixed].translation
        self.fe_r_pos_body[:] = self.data_biped_fixed.oMi[self.r_ankle_joint_fixed].translation
        self.fe_l_rot_body[:, :] = self.data_biped_fixed.oMi[self.l_ankle_joint_fixed].rotation
        self.fe_r_rot_body[:, :] = self.data_biped_fixed.oMi[self.r_ankle_joint_fixed].rotation
        self.hip_l_pos_body[:] = self.data_biped_fixed.oMi[self.l_hip_joint_fixed].translation
        self.hip_r_pos_body[:] = self.data_biped_fixed.oMi[self.r_hip_joint_fixed].translation

    def intQuat(self, quat: sm.UnitQuaternion, w: np.ndarray) -> sm.UnitQuaternion:
        Rcur = quat.R
        Rinc = np.eye(3)
        theta = np.linalg.norm(w)
        if theta > 1e-8:
            w_norm = w / theta
            a = np.array([[0, -w_norm[2], w_norm[1]],
                          [w_norm[2], 0, -w_norm[0]],
                          [-w_norm[1], w_norm[0], 0]])
            Rinc = np.eye(3) + a * np.sin(theta) + a @ a * (1 - np.cos(theta))
        Rend = sm.SO3(Rcur @ Rinc)
        quatRes = Rend.UnitQuaternion()
        return quatRes

    def integrateDIY(self, qI: np.ndarray, dqI: np.ndarray) -> np.ndarray:
        qRes = np.zeros(self.model_nv + 1)
        wBes = np.array([dqI[3], dqI[4], dqI[5]])
        quatNow = sm.UnitQuaternion([qI[6], qI[3], qI[4], qI[5]])
        quatNew = self.intQuat(quatNow, wBes)
        qRes[:] = qI
        qRes[0] += dqI[0]
        qRes[1] += dqI[1]
        qRes[2] += dqI[2]
        qRes[3] = quatNew.A[1]
        qRes[4] = quatNew.A[2]
        qRes[5] = quatNew.A[3]
        qRes[6] = quatNew.A[0]

        for i in range(self.model_nv - 6):
            qRes[7 + i] += dqI[6 + i]
        return qRes

    def computeDyn(self):
        pin.crba(self.model_biped, self.data_biped, self.q)
        self.dyn_M[:, :] = self.data_biped.M

        pin.computeMinverse(self.model_biped, self.data_biped, self.q)
        self.dyn_M_inv[:, :] = self.data_biped.Minv

        pin.computeCoriolisMatrix(self.model_biped, self.data_biped, self.q, self.dq)
        self.dyn_C[:, :] = self.data_biped.C

        pin.computeGeneralizedGravity(self.model_biped, self.data_biped, self.q)
        self.dyn_G[:] = self.data_biped.g

        pin.dccrba(self.model_biped, self.data_biped, self.q, self.dq)
        pin.computeCentroidalMomentum(self.model_biped, self.data_biped, self.q, self.dq)
        self.dyn_Ag[:, :] = self.data_biped.Ag[:, :]
        self.dyn_dAg[:, :] = self.data_biped.dAg[:, :]

        self.dyn_Non[:] = self.dyn_C @ self.dq + self.dyn_G

        pin.ccrba(self.model_biped, self.data_biped, self.q, self.dq)
        self.inertia[:, :] = self.data_biped.Ig.inertia

        self.CoM_pos[:] = self.data_biped.com[0]

        Mpj = np.eye(self.model_nv)
        Mpj_inv = np.eye(self.model_nv)
        Mpj[0: 3, 0: 3] = self.base_rot.T
        Mpj[3: 6, 3: 6] = self.base_rot.T
        Mpj_inv[0: 3, 0: 3] = self.base_rot
        Mpj_inv[3: 6, 3: 6] = self.base_rot
        self.dyn_M[:, :] = Mpj_inv @ self.dyn_M @ Mpj
        self.dyn_M_inv[:, :] = Mpj_inv @ self.dyn_M_inv @ Mpj
        self.dyn_C[:, :] = Mpj_inv @ self.dyn_C @ Mpj
        self.dyn_G[:] = Mpj_inv @ self.dyn_G
        self.dyn_Non[:] = Mpj_inv @ self.dyn_Non

    def computeInk_Leg(self, Rdes_L: np.ndarray, Pdes_L: np.ndarray, Rdes_R: np.ndarray, Pdes_R: np.ndarray) -> IkRes:
        oMdesL = pin.SE3(Rdes_L, Pdes_L)
        oMdesR = pin.SE3(Rdes_R, Pdes_R)
        q_ik = np.zeros(self.model_biped_fixed.nv)
        q_ik[3] = 1.0
        q_ik[9] = 1.0

        eps = 1e-4
        IT_MAX = 100
        DT = 6e-1
        damp = 1e-2
        JL = np.zeros((6, self.model_biped_fixed.nv))
        JR = np.zeros((6, self.model_biped_fixed.nv))
        JCompact = np.zeros((12, self.model_biped_fixed.nv))

        success = False
        errL = np.zeros(6)
        errR = np.zeros(6)
        errCompact = np.zeros(12)
        v = np.zeros(self.model_biped_fixed.nv)

        J_Idx_l = self.l_ankle_joint_fixed
        J_Idx_r = self.r_ankle_joint_fixed
        itr_count = 0
        for itr_count in range(IT_MAX):
            pin.forwardKinematics(self.model_biped_fixed, self.data_biped_fixed, q_ik)
            iMdL = self.data_biped_fixed.oMi[J_Idx_l].actInv(oMdesL)
            iMdR = self.data_biped_fixed.oMi[J_Idx_r].actInv(oMdesR)
            errL = pin.log6(iMdL).vector
            errR = pin.log6(iMdR).vector
            errCompact[0: 6] = errL
            errCompact[6: 12] = errR
            if np.linalg.norm(errCompact) < eps:
                success = True
                break

            JL = pin.computeJointJacobian(self.model_biped_fixed, self.data_biped_fixed, q_ik, J_Idx_l)
            JR = pin.computeJointJacobian(self.model_biped_fixed, self.data_biped_fixed, q_ik, J_Idx_r)
            W = np.eye(self.model_biped_fixed.nv, self.model_biped_fixed.nv)

            JL[0: 6, 16: 19] = 0
            JR[0: 6, 16: 19] = 0

            JlogL = pin.Jlog6(iMdL.inverse())
            JlogR = pin.Jlog6(iMdR.inverse())

            JL = -JlogL @ JL
            JR = -JlogR @ JR

            JCompact[0: 6, 0: self.model_biped_fixed.nv] = JL
            JCompact[6: 12, 0: self.model_biped_fixed.nv] = JR
            JJt = np.zeros((12, 12))
            JJt = JCompact @ W @ JCompact.T
            JJt = JJt + damp * np.eye(12)
            v = -W @ JCompact.T @ np.linalg.pinv(JJt) @ errCompact
            q_ik = pin.integrate(self.model_biped_fixed, q_ik, v * DT)
        res = IkRes()
        res.err = errCompact
        res.itr = itr_count

        if success:
            res.status = 0
        else:
            res.status = -1

        res.jointPosRes = q_ik.copy()
        return res
